import os
import random
import numpy as np
import tensorflow as tf
from PIL import Image

number = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
alphabet = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']
ALPHABET = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z']
CHAR_SET = number + alphabet + ALPHABET
CHAR_SET_LEN = len(CHAR_SET)
MAX_CAPTCHA = 4
IMAGE_HEIGHT = 26
IMAGE_WIDTH= 85
SAVE_PATH = "D:\\PythonProject\\vc_recognition_test\\"
BASE_PATH = "D:\\PythonProject\\vc_recognition_test\\vc_base\\"

def get_name_and_Image():
    all_image = os.listdir(BASE_PATH)
    random_file = random.randint(0,len(all_image)-1)
    text = os.path.splitext(all_image[random_file])[0]
    image = Image.open(BASE_PATH + all_image[random_file])
    image = np.array(image)
    return text,image

def text2vec(text):
    vector = np.zeros([MAX_CAPTCHA,CHAR_SET_LEN])
    for i, c in enumerate(text):
        idx = CHAR_SET.index(c)
        vector[i][idx] = 1.0
    return vector

def vec2text(vec):
    text = []
    for i, c in enumerate(vec):
        text.append(CHAR_SET[c])
    return "".join(text)

def convert2gray(img):
    if len(img.shape) > 2:
        gray = np.mean(img, -1)
        return gray
    else:
        return img

def get_next_batch(batch_size=128):
    batch_x = np.zeros([batch_size, IMAGE_HEIGHT,IMAGE_WIDTH,1])
    batch_y = np.zeros([batch_size, MAX_CAPTCHA,CHAR_SET_LEN])

    def wrap_gen_captcha_text_and_image():
        while True:
            text, image = get_name_and_Image()
            if image.shape == (IMAGE_HEIGHT, IMAGE_WIDTH, 3):
                return text, image
    
    for i in range(batch_size):
        text,image = wrap_gen_captcha_text_and_image()
        image = convert2gray(image)
        batch_x[i, : , : ,0] = image
        batch_y[i, : , : ] = text2vec(text)
    return batch_x,batch_y

def crack_captcha_cnn():
    model = tf.keras.Sequential()

    model.add(tf.keras.layers.Conv2D(32, (3, 3),padding="same"))
    model.add(tf.keras.layers.PReLU())
    model.add(tf.keras.layers.MaxPool2D((2, 2), strides=2))

    model.add(tf.keras.layers.Conv2D(64, (5, 5),padding="same"))
    model.add(tf.keras.layers.PReLU())
    model.add(tf.keras.layers.MaxPool2D((2, 2), strides=2))

    model.add(tf.keras.layers.Conv2D(128, (5, 5),padding="same"))
    model.add(tf.keras.layers.PReLU())
    model.add(tf.keras.layers.MaxPool2D((2, 2), strides=2))

    model.add(tf.keras.layers.Flatten())
    model.add(tf.keras.layers.Dense(MAX_CAPTCHA * CHAR_SET_LEN))
    model.add(tf.keras.layers.Reshape([MAX_CAPTCHA, CHAR_SET_LEN]))
    model.add(tf.keras.layers.Softmax())

    return model

def train_crack_captcha_cnn():
    try:
        model = tf.keras.models.load_model(SAVE_PATH + 'model')
    except Exception as e:
        print('#######Exception', e)
        model = crack_captcha_cnn()

    model.compile(optimizer='Adam', metrics=['accuracy'], loss='categorical_crossentropy')

    for times in range(50):
        batch_x, batch_y = get_next_batch(512)
        print('times=', times, ' batch_x.shape=', batch_x.shape, ' batch_y.shape=', batch_y.shape)
        model.fit(batch_x, batch_y, epochs=4)
        print("y预测=\n", np.argmax(model.predict(batch_x), axis=2))
        print("y实际=\n", np.argmax(batch_y, axis=2))

        if 0 == times % 10:
            print("save model at times=", times)
            model.save(SAVE_PATH + 'model')

def main():
    train_crack_captcha_cnn()

if __name__ == '__main__':
    main()
